### mplScatterplot.  Scatterplot class for matplotLib.

from mplPlot import *
import matplotlib.path as mpath
import matplotlib.patches as mpatches

class mplScatterplot(mplPlot):
    def __init__(self, parent):
        mplPlot.__init__(self, parent)
        self.state = 0 ## the state of the plot 0 == no selections made, 1 == selection in progress, 2 == selection done
        self.pathCollection = [] ## a collection of paths that can be compared with.
        self.interactionMode = 'None'
        self.currentName = 0
        self.currentPath = None
        self.currentPathData = []
        self.workingColor = '#ff0000'
        self.conection = self.canvas.mpl_connect('button_press_event', self.onclick)
        
    def setMode(self, mode):
        self.interactionMode = mode
        
    def setSelectionMode(self, mode = True):
        if mode:
            self.interactionMode = 'selection'
        else:
            self.interactionMode = 'None'
        
        
    def onclick(self, event):
        print event.button, event.x, event.y, event.xdata, event.ydata
        if self.interactionMode == 'selection':
            self.setSelection(event)
            
    def setSelection(self, event):
        if self.state == 0:
            self.currentPathData = []  # make a new current path since this is the first one.
            ## we are starting a new path so 
            self.currentPath = mpath.Path
            self.currentPathData.append((self.currentPath.MOVETO, (event.xdata, event.ydata)))
            self.subPlot.annotate(str(self.currentName), xy = (event.x - 25, event.y), xycoords = 'figure pixels')
            self.subPlot.plot(event.xdata, event.ydata, 'mo')
            self.state = 1
        elif self.state == 1:
            if event.button == 3:
                self.currentPathData.append((self.currentPath.CLOSEPOLY, (event.xdata, event.ydata)))
                codes, verts = zip(*self.currentPathData)
                path = mpath.Path(verts, codes)
                self.pathCollection.append((path, self.workingColor, self.currentName))
                self.currentName += 1
                patch = mpatches.PathPatch(path, facecolor = self.workingColor, edgecolor = 'blue', alpha = 0.5)
                
                self.subPlot.add_patch(patch)
                xData = []
                yData = []
                for v in verts[:-1]:  ## append the verts together to make lines
                    xData.append(v[0])
                    yData.append(v[1])
                xData.append(verts[0][0])  ### make a loop
                yData.append(verts[0][1])
                self.subPlot.add_line(matplotlib.lines.Line2D(xData, yData))
                self.state = 0
            else:
                self.currentPathData.append((self.currentPath.LINETO, (event.xdata, event.ydata)))
                codes, verts = zip(*self.currentPathData)
                print verts, 'Verts'
                xData = []
                yData = []
                for v in verts:
                    xData.append(v[0])
                    yData.append(v[1])
                self.subPlot.add_line(matplotlib.lines.Line2D(xData, yData))
                self.subPlot.plot(event.xdata, event.ydata, 'mo')
        
        self.canvas.draw()
        print self.currentPathData
        
    def setCurveEnabled(self):
        pass
    def closeCurve(self):
        self.currentPathData.append((self.currentPath.CLOSEPOLY, (0, 0)))
        codes, verts = zip(*self.currentPathData)
        path = mpath.Path(verts, codes)
        self.pathCollection.append((path, self.workingColor, self.currentName))
        self.currentName += 1
        patch = mpatches.PathPatch(path, facecolor = self.workingColor, edgecolor = 'blue', alpha = 0.5)
        
        self.subPlot.add_patch(patch)
        xData = []
        yData = []
        for v in verts[:-1]:  ## append the verts together to make lines
            xData.append(v[0])
            yData.append(v[1])
        xData.append(verts[0][0])  ### make a loop
        yData.append(verts[0][1])
        self.subPlot.add_line(matplotlib.lines.Line2D(xData, yData))
        self.state = 0
    def clearCurve(self):
        self.state = 0 ## the state of the plot 0 == no selections made, 1 == selection in progress, 2 == selection done
        self.pathCollection = [] ## a collection of paths that can be compared with.
        self.currentPath = None
        self.currentPathData = []
    def getSelected(self, x = None, y = None):
        ## return a collection of points that are in each of the paths.  This will normally take data from the current plot put could accept any arbitrary points that the widget supplies.
        if not x:
            x = self.x
        if not y:
            y = self.y
        if len(x) != len(y):
            raise Exception, 'Length X and Y not the same'
        
        r = {}
        pathNumber = 0
        for path, color, name in self.pathCollection:
            r[str(name)] = []
            for i in range(len(x)):
                if path.contains_point([x[i], y[i]]):
                    r[str(name)].append(1)
                else:
                    r[str(name)].append(0)
            pathNumber += 1
        return r
    def makePlot(self, x, y, s=20, c = 'b', marker = 'o', cmap=None, norm = None, vmin = None, vmax = None, alpha = 0.5, linewidths = None, verts = None, resetPaths = True, **kwargs):
        self.subPlot.clear()
        self.x = x
        self.y = y
        self.subPlot.scatter(x = x, y = y, s = s, c = c, marker = marker, cmap = cmap, norm = norm, vmin = vmin, vmax = vmax, alpha = alpha, linewidths = linewidths, verts = verts, **kwargs)
        
        if resetPaths:
            self.pathCollection = []
        else:
            for path, color in self.pathCollection:
                patch = mpatches.PathPatch(path, facecolor = color, edgecolor = 'blue', alpha = 0.5)
                self.subPlot.add_patch(patch)
                
        self.canvas.draw()